/*
 *        (C) COPYRIGHT Ingenic Limited.
 *             ALL RIGHTS RESERVED
 *
 * File       : inference_yolov5s.cpp
 * Authors    : ljhong, cywang
 * Create Time: 2024-11-21:15:57:21
 * Description:
 * 
 */

#define STB_IMAGE_IMPLEMENTATION
#include "./stb/stb_image.h"
#include "./stb/drawing.hpp"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "./stb/stb_image_write.h"
static const uint8_t color[3] = {0xff, 0, 0};

#include "ingenic_mnni.h"
#include "common/common_utils.h"
#include "utils/imgproc.h"
#include "utils/postproc.h"

#include <cmath>
#include <errno.h>
#include <fstream>
#include <iostream>
#include <memory>
#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <cstring>
#include <string.h>
#include <vector>
#include <algorithm>
#include <sys/stat.h>   
#include <sys/types.h>
#include <dirent.h>     
#include <unistd.h>     
#include <sstream>      
#include <iomanip>      
#include <numeric>  // 添加 numeric 头文件，用于 std::inner_product

#define IS_ALIGN_64(x) (((size_t)x) & 0x3F)


int ori_img_w;
int ori_img_h;
float conf_thres = 0.6;
float nms_thres = 0.6;
int classes = 1;
int box_num = 3;

//  yolov5
vector<float> strides_yolov5 = {8.0, 16.0, 32.0};
std::vector<float> anchor = {10, 13, 16, 30, 33, 23, 
                             30, 61, 62, 45, 59, 119, 
                             116, 90, 156, 198, 373, 326};
//  yolov8
std::vector<int> strides_yolov8 = {8,16,32};
int reg_max = 16;

using namespace std;
using namespace magik::uranus;

#ifdef URANUS_NV12
magik::uranus::ChannelLayout DATA_FORMAT = magik::uranus::ChannelLayout::NV12;
#else
magik::uranus::ChannelLayout DATA_FORMAT = magik::uranus::ChannelLayout::RGBA;
#endif

typedef unsigned char uint8_t;
#define RET_CHECKER(ret, msg)                                           \
    do {                                                                \
        if (!(ret)) {                                                   \
            fprintf(stderr, "%s [%s %d]\n", msg, __func__, __LINE__);   \
            return -1;                                                  \
        }                                                               \
    } while(0)

typedef struct {
    uint8_t *data;
    int width;
    int height;
}FrameInfo;

void print_memory_usage_mb(const std::string& tag) {
    std::ifstream meminfo("/proc/meminfo");
    std::string line;
    long mem_total = 0, mem_free = 0, buffers = 0, cached = 0;

    // 从 /proc/meminfo 中提取关键字段（单位KB）
    while (std::getline(meminfo, line)) {
        if (line.find("MemTotal:") == 0) {
            sscanf(line.c_str(), "MemTotal: %ld kB", &mem_total);
        } else if (line.find("MemFree:") == 0) {
            sscanf(line.c_str(), "MemFree: %ld kB", &mem_free);
        } else if (line.find("Buffers:") == 0) {
            sscanf(line.c_str(), "Buffers: %ld kB", &buffers);
        } else if (line.find("Cached:") == 0) {
            sscanf(line.c_str(), "Cached: %ld kB", &cached);
        }
    }

    // 转换为MB并打印
    std::cout << "===== Memory Usage [" << tag << "] (MB) =====" << std::endl;
    // std::cout << "Total RAM: " << mem_total / 1024 << " MB" << std::endl;
    // std::cout << "Free RAM: " << mem_free / 1024 << " MB" << std::endl;
    // std::cout << "Buffers: " << buffers / 1024 << " MB" << std::endl;
    // std::cout << "Cached: " << cached / 1024 << " MB" << std::endl;
    std::cout << "Available (Free + Buffers + Cached): " 
              << (mem_free + buffers + cached) / 1024 << " MB" << std::endl;
    std::cout << "=========================================" << std::endl;
}

float L2_distance(const std::vector<float>& a, const std::vector<float>& b) {
    // 使用 std::inner_product 计算平方差和
    float sum = std::inner_product(
        a.begin(), a.end(), b.begin(), 0.0f,
        std::plus<float>(),
        [](float x, float y) { float d = x - y; return d * d; }
    );

    return std::sqrt(sum);
}

double reduceL2(const float* input, int dim) {
    double squaredSum = 0.0;
    for (int i = 0; i< dim; ++i) {
         squaredSum += input[i] * input[i];
    }
    double l2Norm = std::sqrt(squaredSum);
    return l2Norm;
}


static int fill_dim(uint8_t *data, int h, int w) {
    for (int i = 0; i < h*w; i++) {
        data[i*4+3] = 0;
    }
    return 0;
}

int rgb2bgr(uint8_t *data, int data_cnt){
    uint8_t* p = data;
    int pixel_c = data_cnt / 4;
    for (int i=0; i < pixel_c; ++i) {
        uint8_t t = p[0];
        p[0] = p[2];
        p[2] = t;
        p += 4;
    }
    return 0;
}

void read_img(FrameInfo &in, Img &img, const std::string &img_path) {
    int w, h, c;
    // 仅加载一次4通道RGBA数据（包含Alpha通道）
    unsigned char *rgba_data = stbi_load(img_path.c_str(), &w, &h, &c, 4);
    // RET_CHECKER(rgba_data != nullptr, "stbi_load failed");

    // 填充FrameInfo（直接使用RGBA数据）
    in.data = rgba_data;
    in.width = w;
    in.height = h;
    printf("input===> %d %d \n", in.width, in.height);

    // 从RGBA数据中提取RGB三通道（避免重复加载）
    unsigned char *rgb_data = (unsigned char *)malloc(w * h * 3);
    // RET_CHECKER(rgb_data != nullptr, "malloc for rgb_data failed");
    for (int i = 0; i < w * h; i++) {
        rgb_data[i*3 + 0] = rgba_data[i*4 + 0]; // R
        rgb_data[i*3 + 1] = rgba_data[i*4 + 1]; // G
        rgb_data[i*3 + 2] = rgba_data[i*4 + 2]; // B
    }

    // 填充Img结构体
    img.w = w;
    img.h = h;
    img.c = 3;
    img.w_stride = w * 3;
    img.data = rgb_data;

    fill_dim(in.data, w, h); // 保留Alpha通道处理逻辑
}

void prepare_data(FrameInfo &in, magik::uranus::Tensor &out, magik::uranus::Box_t tmp_box, int input_w, int input_h) {
    int in_w = in.width;
    int in_h = in.height;
    printf("perpare_data, in_h in_w %d %d\n", in.height, in.width);
	magik::uranus::shape_t in_shape;
	magik::uranus::DataFormat data_format;
    if (DATA_FORMAT == magik::uranus::ChannelLayout::NV12) {
        in_shape = {1, in_h, in_w, 1};
		data_format = magik::uranus::DataFormat::NV12;
        
    } else {
        in_shape = {1, in_h, in_w, 4};
		data_format = magik::uranus::DataFormat::NHWC;
    }
	auto in_tensor = magik::uranus::Tensor(in_shape, magik::uranus::DataType::UINT8, data_format);
	uint8_t *in_tensor_data = in_tensor.mudata<uint8_t>();
	memcpy((void *)in_tensor_data, (void *)in.data, in_tensor.get_bytes_size());

    // resize
    magik::uranus::BsCommonParam param;
    param.pad_val = 114;
    param.pad_type = magik::uranus::BsPadType::SYMMETRY;
    param.in_layout = DATA_FORMAT;
    param.out_layout = magik::uranus::ChannelLayout::RGBA;
    param.input_height = in_h;
    param.input_width = in_w;
    param.input_line_stride = (DATA_FORMAT == magik::uranus::ChannelLayout::NV12) ? in_w : in_w * 4;
    param.addr_attr.vir_addr = in_tensor_data;
    vector<magik::uranus::Box_t> output_tmp;
    // magik::uranus::Box_t tmp_box{0, 0, static_cast<float>(in_w), static_cast<float>(in_h)};
    output_tmp.push_back(tmp_box);
    printf("work_h:%d work_w:%d \n", input_h, input_w);
    out.reshape({1, input_h, input_w, 4});
    uint8_t *out_data = out.mudata<uint8_t>();
    vector<magik::uranus::Tensor> output_tesnor;
    output_tesnor.push_back(out);
    magik::uranus::crop_common_resize(output_tesnor, output_tmp, magik::uranus::AddressLocate::NMEM_VIRTUAL, &param);

    in_tensor.free_data();

    if (IS_ALIGN_64(out_data) != 0) {
        fprintf(stderr, "input addr not align to 128 bytes.\n");
        return ;
    }
}

int frame_to_img(const FrameInfo &frame, Img &img) {
    // 检查输入有效性
    if (frame.data == nullptr || frame.width <= 0 || frame.height <= 0) {
        fprintf(stderr, "Invalid FrameInfo input\n");
        return -1;
    }

    // 分配RGB图像内存
    unsigned char *imagedata = (unsigned char *)malloc(frame.width * frame.height * 3);
    if (!imagedata) {
        fprintf(stderr, "Memory allocation failed\n");
        return -1;
    }

    // 从RGBA转换为RGB
    for (int i = 0; i < frame.height; i++) {
        for (int j = 0; j < frame.width; j++) {
            int src_idx = (i * frame.width + j) * 4;  // RGBA索引
            int dst_idx = (i * frame.width + j) * 3;  // RGB索引
            
            // 复制RGB通道，忽略Alpha通道
            imagedata[dst_idx + 0] = frame.data[src_idx + 0]; // R
            imagedata[dst_idx + 1] = frame.data[src_idx + 1]; // G
            imagedata[dst_idx + 2] = frame.data[src_idx + 2]; // B
        }
    }

    // 填充Img结构体
    img.w = frame.width;
    img.h = frame.height;
    img.c = 3;
    img.w_stride = frame.width * 3;
    img.data = imagedata;

    return 0;
}


int tensor_to_img(const magik::uranus::Tensor &tensor, Img &img) {
    // 检查Tensor有效性
    auto shape = tensor.shape();
    if (shape.size() != 4 || shape[0] != 1) {  // 假设Tensor格式为NHWC（1, H, W, C）
        fprintf(stderr, "Invalid tensor shape for image conversion\n");
        return -1;
    }
    int height = shape[1];
    int width = shape[2];
    int channels = shape[3];
    if (channels < 3) {  // 至少需要RGB三通道
        fprintf(stderr, "Tensor channel count must be >=3 for RGB conversion\n");
        return -1;
    }

    // 获取Tensor数据指针（假设是uint8_t类型）
    const uint8_t *tensor_data = tensor.vdata<uint8_t>();
    if (!tensor_data) {
        fprintf(stderr, "Tensor data pointer is null\n");
        return -1;
    }

    // 分配RGB图像内存（忽略Alpha通道）
    unsigned char *imagedata = (unsigned char *)malloc(width * height * 3);
    if (!imagedata) {
        fprintf(stderr, "Memory allocation failed for crop_img\n");
        return -1;
    }

    // 从Tensor（RGBA）转换为RGB
    for (int i = 0; i < height; i++) {
        for (int j = 0; j < width; j++) {
            int src_idx = (i * width + j) * channels;  // Tensor数据索引（RGBA）
            int dst_idx = (i * width + j) * 3;         // RGB目标索引
            
            // 复制RGB通道，忽略Alpha通道（如果存在）
            imagedata[dst_idx + 0] = tensor_data[src_idx + 0]; // R
            imagedata[dst_idx + 1] = tensor_data[src_idx + 1]; // G
            imagedata[dst_idx + 2] = tensor_data[src_idx + 2]; // B
        }
    }
    // 填充Img结构体
    img.w = width;
    img.h = height;
    img.c = 3;
    img.w_stride = width * 3;
    img.data = imagedata;

    return 0;
}


void genBBox_yolov5(std::vector<magik::uranus::Tensor> out_res, std::vector<magik::uranus::ObjBbox_t>& candidate_boxes, int img_w, int img_h) {
    DetectorType_t detector_type = DetectorType_t::YOLOV5;
    magik::uranus::generate_box(out_res, strides_yolov5, anchor, candidate_boxes, img_w, img_h, classes, box_num, conf_thres, detector_type);
}


void genBBox_yolov8(std::vector<magik::uranus::Tensor> out_res, std::vector<magik::uranus::ObjBbox_t>& candidate_boxes) {   
    magik::uranus::generateBBox_yolov8(out_res, candidate_boxes, strides_yolov8, conf_thres, classes, reg_max);
}

void scale_boxes(int ori_w, int ori_h, int new_w, int new_h, std::vector<magik::uranus::ObjBbox_t> &out_boxes) {
    float scale;
    float pad_x, pad_y;
    scale = std::min((float(new_w) / ori_w), (float(new_h) / ori_h));
    pad_x = (new_w - ori_w * scale) / 2;
    pad_y = (new_h - ori_h * scale) / 2;
    for (unsigned int i = 0; i < out_boxes.size(); i++)
    {
        out_boxes[i].box.x0 = round(std::max((out_boxes[i].box.x0 - pad_x) / scale, float(0)));
        out_boxes[i].box.x1 = round(std::min((out_boxes[i].box.x1 - pad_x) / scale, float(ori_w - 1)));
        out_boxes[i].box.y0 = round(std::max((out_boxes[i].box.y0 - pad_y) / scale, float(0)));
        out_boxes[i].box.y1 = round(std::min((out_boxes[i].box.y1 - pad_y) / scale, float(ori_h - 1)));
    }
}

void mobilenet_run_model(
    magik::uranus::Network *model,
    magik::uranus::Tensor &img,
    std::vector<float> &out,  // 改为引用传递
    int size
) {

    int img_size = img.get_bytes_size();
    auto input_names = model->get_input_names();
    // printf("input_names: %s\n", input_names[0].c_str());
    for (auto n : input_names) {
        auto input = model->get_input(n);
        uint8_t *indata = (uint8_t*)(input.vdata<uint8_t>());
        memcpy((void *)indata, (void *)img.vdata<uint8_t>(), img_size * sizeof(uint8_t));

        // 前向推理
        int RUN_CNT = 1;
        for (int i = 0; i < RUN_CNT; i++) {
            model->run();
        }

        // 获取输出
        auto output_names = model->get_output_names();
        for (auto n : output_names) {
            // std::cout << "output_names: " << n << std::endl;
            auto output = model->get_output(n);
            const float* r_ptr1 = output.vdata<float>();
            double l2_data = reduceL2(r_ptr1, size);

            // 将结果写入 out（引用传递，外部可见）
            for (int i = 0; i < size; ++i) {
                out[i] = r_ptr1[i] / l2_data;
            }
            auto shape = output.shape();
            // std::cout<<"mobileNet output_shape: "<< shape[0] << " " <<std::endl;
        }
    }
}

void yolov5_run_model(magik::uranus::Network *model, magik::uranus::Tensor &img, std::vector<magik::uranus::Tensor> &outputs) {
    int img_size = img.get_bytes_size();
    auto input_names = model->get_input_names();
    for(auto n: input_names) {
        std::cout<<"input_names: "<< n <<std::endl;
        auto input = model->get_input(n);
        auto shape = input.shape();
        // std::cout<<"input_shape: "<< shape[0] << " " << shape[1] << " " << shape[2] << " " << shape[3] << " " <<std::endl;

        uint8_t *indata = (uint8_t*)(input.vdata<uint8_t>());
        memcpy((void *)indata, (void *)img.vdata<uint8_t>(), img_size * sizeof(uint8_t));

        // const unsigned char* in_ptr = input.vdata<uint8_t>();
        // FILE *fp = fopen(std::string(n+".res").c_str(), "wb");
        // if (fp) {
        //     fwrite(in_ptr, 1, (input.get_bytes_size()), fp);
        //     fclose(fp);
        // } else {
        //     printf("open %s file failed\n", std::string(n+".res").c_str());
        //     return;
        // }
    }

    // net forword
    int RUN_CNT = 1;
    for (int i = 0; i < RUN_CNT; i++) {
        model->run();
    }

    // get outputs
    auto output_names = model->get_output_names();
    for(auto n: output_names) {
        // std::cout<<"output_names: "<< n <<std::endl;
        auto output = model->get_output(n);
        auto shape = output.shape();
        // std::cout<<"output_shape: "<< shape[0] << " " << shape[1] << " " << shape[2] << " " << shape[3] << " " <<std::endl;

        outputs.push_back(output);

        // const unsigned char* out_ptr = output.vdata<uint8_t>();
        // FILE *fp = fopen(std::string(n+".res").c_str(), "wb");
        // if (fp) {
        //     fwrite(out_ptr, 1, (output.get_bytes_size()), fp);
        //     fclose(fp);
        // } else {
        //     printf("open %s file failed\n", std::string(n+".res").c_str());
        //     return;
        // }
    }

    return;
}


void yolov8_run_model(magik::uranus::Network *model, magik::uranus::Tensor &img, std::vector<magik::uranus::Tensor> &outputs) {
    int img_size = img.get_bytes_size();
    auto input_names = model->get_input_names();
    for(auto n: input_names) {
        std::cout<<"input_names: "<< n <<std::endl;
        auto input = model->get_input(n);
        uint8_t *indata = (uint8_t*)(input.vdata<uint8_t>());
        memcpy((void *)indata, (void *)img.vdata<uint8_t>(), img_size * sizeof(uint8_t));

        // const unsigned char* in_ptr = input.vdata<uint8_t>();
        // FILE *fp = fopen(std::string(n+".res").c_str(), "wb");
        // if (fp) {
        //     fwrite(in_ptr, 1, (input.get_bytes_size()), fp);
        //     fclose(fp);
        // } else {
        //     printf("open %s file failed\n", std::string(n+".res").c_str());
        //     return;
        // }
    }

    // net forword
    int RUN_CNT = 1;
    for (int i = 0; i < RUN_CNT; i++) {
        model->run();
    }

    // get outputs
    auto output_names = model->get_output_names();
    for(auto n: output_names) {
        std::cout<<"output_names: "<< n <<std::endl;
        auto output = model->get_output(n);
        outputs.push_back(output);

        // const unsigned char* out_ptr = output.vdata<uint8_t>();
        // FILE *fp = fopen(std::string(n+".res").c_str(), "wb");
        // if (fp) {
        //     fwrite(out_ptr, 1, (output.get_bytes_size()), fp);
        //     fclose(fp);
        // } else {
        //     printf("open %s file failed\n", std::string(n+".res").c_str());
        //     return;
        // }
    }

    // std::string cls_stride_8 = "851_last_layer";
    // std::string cls_stride_16 = "874_last_layer";
    // std::string cls_stride_32 = "897_last_layer";
    // std::string bbox_stride_8 = "840_last_layer";
    // std::string bbox_stride_16 = "863_last_layer";
    // std::string bbox_stride_32 = "886_last_layer";

    // auto cls_stride_8_output = model->get_output(cls_stride_8);
    // auto bbox_stride_8_output = model->get_output(bbox_stride_8);
    // auto cls_stride_16_output = model->get_output(cls_stride_16);
    // auto bbox_stride_16_output = model->get_output(bbox_stride_16);
    // auto cls_stride_32_output = model->get_output(cls_stride_32);
    // auto bbox_stride_32_output = model->get_output(bbox_stride_32);

    // // Push the outputs to the 'outputs' vector in the order of classification heads followed by detection heads
    // // Note 1: Prioritize classification heads before detection heads
    // // Note 2: The order of elements corresponds to their respective strides
    // outputs.push_back(cls_stride_8_output);
    // outputs.push_back(bbox_stride_8_output);
    // outputs.push_back(cls_stride_16_output);
    // outputs.push_back(bbox_stride_16_output);
    // outputs.push_back(cls_stride_32_output);
    // outputs.push_back(bbox_stride_32_output);

    return;
}

// 在run_detect函数前添加辅助函数：获取文件夹中所有图片路径
std::vector<std::string> get_image_files(const std::string& folder_path) {
    std::vector<std::string> image_files;
    DIR* dir = opendir(folder_path.c_str());
    if (!dir) {
        fprintf(stderr, "Failed to open directory: %s\n", folder_path.c_str());
        return image_files;
    }

    struct dirent* entry;
    while ((entry = readdir(dir)) != nullptr) {
        if (entry->d_type != DT_REG) continue; // 只处理普通文件
        
        std::string filename = entry->d_name;
        // 检查文件扩展名是否为常见图片格式
        if (filename.size() >= 4) {
            std::string ext = filename.substr(filename.size() - 4);
            std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
            if (ext == ".jpg" || ext == ".png" || ext == ".bmp" || ext == ".jpeg") {
                image_files.push_back(folder_path + "/" + filename);
            }
        }
    }
    closedir(dir);
    return image_files;
}

// 修改run_detect函数，使其接收单张图片路径并返回处理结果
// 修改后的函数声明（新增模型指针参数）
// 修改函数声明，新增日志文件流参数
int run_detect_single_image(
    std::string yolo_model_type,
    std::string img_path, int img_w, int img_h, 
    magik::uranus::Network *det_model, 
    magik::uranus::Network *recog_model,
    std::ofstream& undetected_file) {  // 新增参数
    //  /******************load images for inference***********************/
    print_memory_usage_mb("START");
    std::cout << "load image: " << img_path << std::endl;
    FrameInfo input_info = {nullptr, 0, 0};  // 初始化避免野指针
    // 修正：Img 结构体的整数成员用0初始化，指针成员用nullptr
    Img img = {0, 0, 0, 0, nullptr};   

    read_img(input_info, img, img_path.c_str());
    print_memory_usage_mb("AFTER_READ_IMG"); 
    // 检查图片加载是否成功（input_info.data可能为空）

    // 检查图片尺寸是否超过阈值（2000 * 2000）
    if (input_info.width * input_info.height > 2000 * 2000) {
        std::cerr << "警告：图片尺寸过大（" << input_info.width << "x" << input_info.height 
                  << "），超过2000*2000阈值，跳过处理" << std::endl;
        free(input_info.data);    // 释放RGBA数据
        input_info.data = nullptr;
        free(img.data);           // 释放RGB数据
        img.data = nullptr;
        return 0;  
    }

    magik::uranus::Box_t tmp_box;
    tmp_box = {0, 0, static_cast<float>(input_info.width), static_cast<float>(input_info.height)};
    const int yolo_input_w = 640;
    const int yolo_input_h = 640;
    /******************prepare data***********************/
    magik::uranus::shape_t yolo_working_img_shape = {1, yolo_input_h, yolo_input_w, 4};
    magik::uranus::Tensor yolo_working_img(yolo_working_img_shape);
    prepare_data(input_info,  yolo_working_img, tmp_box, yolo_input_w, yolo_input_h);
    print_memory_usage_mb("AFTER_PREPARE_DATA YOLO"); // 预处理后内存                        
    /******************model forward***********************/
    std::vector<magik::uranus::Tensor> outputs;
    if (yolo_model_type == "yolov5") {
        yolov5_run_model(det_model, yolo_working_img, outputs);  // 直接使用传入的det_model
        yolo_working_img.free_data();
    } else if (yolo_model_type == "yolov8") {
        yolov8_run_model(det_model, yolo_working_img, outputs);  // 直接使用传入的det_model
        yolo_working_img.free_data();
    }
    else {
        std::cerr << "错误：不支持的yolo模型类型" << std::endl;
        return -1;
    }

    /******************postprocess***********************/
    std::vector<magik::uranus::ObjBbox_t> temp_boxes;
    vector<magik::uranus::ObjectBbox_t> output_boxes;
    if (yolo_model_type == "yolov5") {
        genBBox_yolov5(outputs, temp_boxes, yolo_input_w, yolo_input_h);
    } else if (yolo_model_type == "yolov8") {
        genBBox_yolov8(outputs, temp_boxes);
    }
    magik::uranus::nms(temp_boxes, output_boxes, nms_thres, magik::uranus::NmsType_t::HARD_NMS);
    scale_boxes(input_info.width, input_info.height, yolo_input_w, yolo_input_h, output_boxes);
    int num_boxes = output_boxes.size();
    std::cout << "num_boxes: " << num_boxes << std::endl;

    // 新增：检测到0个框时提前终止，避免后续越界访问
    if (num_boxes == 0) {
        std::cerr << "警告：未检测到有效目标框，跳过当前图片处理" << std::endl;
        
        // 记录未检测图像路径到txt文件
        undetected_file << img_path << std::endl;

        // 保存未检测的原始图片（与检测成功图片同目录）
        if (DATA_FORMAT == magik::uranus::ChannelLayout::RGBA && img.data != nullptr) {
            // 构造输出路径（原文件名_undetected.扩展名）
            std::string img_filename = img_path.substr(img_path.find_last_of("/\\") + 1);
            size_t dot_pos = img_filename.find_last_of('.');
            std::string name_without_ext = (dot_pos == std::string::npos) ? img_filename : img_filename.substr(0, dot_pos);
            std::string ext = (dot_pos == std::string::npos) ? "" : img_filename.substr(dot_pos);
            std::string output_path = "result/" + name_without_ext + "_undetected" + ext;
            
            // 保存JPG（质量60）
            stbi_write_jpg(output_path.c_str(), img.w, img.h, 3, img.data, 60);
            std::cout << "未检测图片已保存: " << output_path << std::endl;
        }

        // 释放已分配的内存
        free(input_info.data);
        input_info.data = nullptr; // 避免悬空指针
        free(img.data);
        return 0;  
    }
    
    // 初始化最大得分和对应索引
    float max_score = 0.0f;
    int max_index = -1;
    
    for (int i = 0; i < int(output_boxes.size()); i++) {
        auto output_box = output_boxes[i];
        printf("box:   ");
        printf("x0: %d ", (int)output_box.box.x0);
        printf("y0: %d ", (int)output_box.box.y0);
        printf("x1: %d ", (int)output_box.box.x1);
        printf("y1: %d ", (int)output_box.box.y1);
        printf("score: %.2f ", output_box.score);
        printf("class: %d ", output_box.class_id);
        printf("\n");
    
        // 跟踪最大得分及其索引
        if (output_box.score > max_score) {
            max_score = output_box.score;
            max_index = i;
        }
    
        Point pt1 = {
            .x = (int)output_box.box.x0,
            .y = (int)output_box.box.y0};
        Point pt2 = {
            .x = (int)output_box.box.x1,
            .y = (int)output_box.box.y1};
        if (DATA_FORMAT == magik::uranus::ChannelLayout::RGBA) {
            sample_draw_box_for_image(&img, pt1, pt2, color, 2);
        }
    }
    

    magik::uranus::Box_t face_box = output_boxes[max_index].box;
    auto make_even = [](float val) {
        int int_val = static_cast<int>(val);
        return (int_val % 2 == 0) ? val : val + 1;
    };
    tmp_box = {
        make_even(face_box.x0),
        make_even(face_box.y0),
        make_even(face_box.x1),
        make_even(face_box.y1)
    };
    /******************MobileNet***********************/
    /******************prepare data***********************/
    const int mobilenet_input_w = 224;
    const int mobilenet_input_h = 224;
    magik::uranus::shape_t recog_working_img_shape = {1, mobilenet_input_h, mobilenet_input_w, 4};
    magik::uranus::Tensor recog_working_img(recog_working_img_shape);
    
    prepare_data(input_info,  recog_working_img, tmp_box, mobilenet_input_w, mobilenet_input_h);
    print_memory_usage_mb("AFTER_PREPARE_DATA MOBILENET"); // 预处理后内存                        
    free(input_info.data);
    input_info.data = nullptr; // 避免悬空指针



    /******************model forward***********************/
    std::vector<float> out1 (128,0);
    mobilenet_run_model(recog_model, recog_working_img, out1, 128);

    /******************save images***********************/
    // 打印最大得分的索引和分数
    if (max_index != -1) {
        printf("Max score index: %d, Max score: %.2f\n", max_index, max_score);
    } else {
        printf("No valid boxes found\n");
    }
    
    if (DATA_FORMAT == magik::uranus::ChannelLayout::RGBA) {
        // 提取图像文件名（不含路径）
        std::string img_filename = img_path.substr(img_path.find_last_of("/\\") + 1);
        
        // 格式化分数为字符串（保留两位小数）
        std::stringstream ss;
        ss << std::fixed << std::setprecision(2) << max_score;
        std::string score_str = ss.str();
        
        // 分割文件名和扩展名
        size_t dot_pos = img_filename.find_last_of('.');
        std::string name_without_ext = img_filename.substr(0, dot_pos);
        std::string ext = img_filename.substr(dot_pos);
        
        // 构建输出路径和文件名：result/原文件名_分数.扩展名（原图）
        std::string output_path = "result/" + name_without_ext + "_score_" + score_str + ext;
        stbi_write_jpg(output_path.c_str(), img.w, img.h, 3, img.data, 60); // 保存原图
        free(img.data);

        // 新增：保存裁剪后的crop_img（添加_crop后缀）
        // Img crop_img;
        // int ret = tensor_to_img(recog_working_img, crop_img);
        // RET_CHECKER(ret == 0, "tensor_to_img conversion failed");
        // std::string crop_output_path = "result/" + name_without_ext + "_score_" + score_str + "_crop" + ext;
        // if (crop_img.data != nullptr) {  // 检查数据有效性
        //     stbi_write_jpg(crop_output_path.c_str(), crop_img.w, crop_img.h, 3, crop_img.data, 60); // 保存裁剪图
        //     free(crop_img.data);  // 释放frame_to_img中malloc的内存
        //     std::cout << "crop_img save to: " << crop_output_path << std::endl;
        // } else {
        //     std::cerr << "crop_img is null, not saved" << std::endl;
        // }
    }
    
        
    /******************save data***********************/
    // 构造与图片命名一致的txt路径（原文件名_分数.txt）
    // std::string img_filename = img_path.substr(img_path.find_last_of("/\\") + 1);
    // std::stringstream ss;
    // ss << std::fixed << std::setprecision(2) << max_score;
    // std::string score_str = ss.str();
    // size_t dot_pos = img_filename.find_last_of('.');
    // std::string name_without_ext = img_filename.substr(0, dot_pos);
    // std::string txt_output_path = "result/" + name_without_ext + "_score_" + score_str + ".txt";

    // std::ofstream out_file(txt_output_path);
    // if (!out_file.is_open()) {
    //     std::cerr << "open file " << txt_output_path << " failed" << std::endl;
    // } else {
    //     for (float value : out1) {
    //         out_file << value << "\n";  // 修正为转义换行符
    //     }
    //     out_file.close();
    //     std::cout << "vector save to: " << txt_output_path << std::endl;
    // }

    // /******************release memory***********************/
    recog_working_img.free_data();
    print_memory_usage_mb("END");
    return 1;
}

int main(int argc, char **argv) {
    std::string yolo_model_type = argv[1];
    std::string face_det_model_path = argv[2];
    std::string recog_model_path = argv[3];
    std::string img_folder = argv[4];

    if (argc != 5)
    {
        printf("%s yolo_model_type face_det_model_path recog_model_path image_folder  \n", argv[0]);
        exit(0);
    }
    
    // 每次运行前删除并重建result文件夹
    system("rm -rf result");
    struct stat info;
    if (stat("result", &info) != 0) {
        if (mkdir("result", 0777) != 0) {
            fprintf(stderr, "Failed to create result directory: %s\n", strerror(errno));
            return -1;
        }
    } else if (!(info.st_mode & S_IFDIR)) {
        fprintf(stderr, "result is not a directory\n");
        return -1;
    }

    /******************load models（提取到外部）***********************/
    Network **derived_models;
    int model_num = 0;
    std::cout << "load face det model: " << face_det_model_path << std::endl;
    if (magik_load_model(derived_models, model_num, face_det_model_path.c_str())) {
        std::cerr << "create model failed." << std::endl;
        exit(1);
    }
    auto det_model = derived_models[0];
    int ret1 = det_model->get_forward_memory_size();
    std::cout << "model forward memory size: " << ret1 << " = " << ret1 / 1024.0 << "KB" << " = " << ret1 / 1024.0 / 1024.0 << "MB" << std::endl;

    Network **recog_derived_models;
    int recog_model_num = 0;
    std::cout << "load recog model: " << recog_model_path << std::endl;
    if (magik_load_model(recog_derived_models, recog_model_num, recog_model_path.c_str())) {
        std::cerr << "create model failed." << std::endl;
        exit(1);
    }
    auto recog_model = recog_derived_models[0];
    int ret2 = recog_model->get_forward_memory_size();
    std::cout << "recog model forward memory size: " << ret2 << " = " << ret2 / 1024.0 << "KB" << " = " << ret2 / 1024.0 / 1024.0 << "MB" << std::endl;

    /******************获取图片列表***********************/
    std::vector<std::string> image_files = get_image_files(img_folder);
    int total_files = image_files.size();
    if (total_files == 0) {
        fprintf(stderr, "No image files found in folder: %s\n", img_folder.c_str());
        return -1;
    }
    printf("Found %d image files, starting inference...\n", total_files);
    
    ori_img_w = -1;
    ori_img_h = -1;

    /******************新增：统计变量初始化***********************/
    int total_processed = 0;    // 总处理图片数
    int detected_count = 0;     // 识别到目标框的图片数

    /******************初始化未检测图像日志文件***********************/
    std::ofstream undetected_file("result/undetected_images.txt");
    if (!undetected_file.is_open()) {
        fprintf(stderr, "Failed to create undetected images log file: %s\n", strerror(errno));
        return -1;
    }

    /******************遍历处理图片（传递模型参数和日志文件）***********************/
    for (int i = 0; i < total_files; ++i) {
        printf("\rProcessing %d/%d: %s\n", i+1, total_files, image_files[i].c_str());
        fflush(stdout);
        // 调用处理函数并传递日志文件引用
        int process_result = run_detect_single_image(
            yolo_model_type, image_files[i], ori_img_w, ori_img_h, 
            det_model, recog_model, undetected_file  // 新增参数
        );
        total_processed++;
        if (process_result == 1) {
            detected_count++;
        }
    }

    /******************关闭日志文件***********************/
    undetected_file.close();

    /******************统一销毁模型（原内部销毁移到此处）***********************/
    magik_destroy_model(det_model);
    magik_destroy_model(recog_model);
    std::cout << "All models destroyed" << std::endl;
    
    /******************新增：输出统计结果***********************/
    printf("\n===== 识别统计结果 =====\n");
    printf("总处理图片数: %d\n", total_processed);
    printf("识别到目标框的图片数: %d\n", detected_count);
    float recognition_rate = (total_processed == 0) ? 0 : (float)detected_count / total_processed * 100;
    printf("识别率: %.2f%%\n", recognition_rate);
    printf("========================\n");

    printf("\nInference completed!\n");
    return 0;
}

